{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tensorflow中的Attention机制源码解析\n",
    "\n",
    "`tensorflow`中的`attention_wrapper`源码解析。\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Attention机制\n",
    "\n",
    "首先，我们回顾一下`attention mechanism`的定义。\n",
    "\n",
    "假设我们的source sequence为长度为$\\text{n}$的序列$\\mathbf{x}$，target sequence为长度为$\\text{m}$的序列$\\mathbf{y}$，则：\n",
    "\n",
    "$$\\mathbf{x}=[x_0,x_1,\\dots,x_n]$$\n",
    "$$\\mathbf{y}=[y_0,y_1,\\dots,y_m]$$\n",
    "\n",
    "在第$i$个time step，source sequence产生一个`hidden state`：\n",
    "\n",
    "$$\\boldsymbol{h}_i = [\\overrightarrow{\\boldsymbol{h}}_i^\\top; \\overleftarrow{\\boldsymbol{h}}_i^\\top]^\\top, i=1,\\dots,n$$\n",
    "\n",
    "其中，$\\overrightarrow{\\boldsymbol{h}}_i^\\top$和$\\overleftarrow{\\boldsymbol{h}}_i^\\top$表示的是`Bidirectional RNN`产生的两个hidden state，如果不是使用BiRNN，则只有一个。\n",
    "\n",
    "在第$t$个time step，decoder产生的`hidden state`表示为：\n",
    "\n",
    "$$\\boldsymbol{s}_t=f(\\boldsymbol{s}_{t-1}, y_{t-1}, \\mathbf{c}_t),t=1,\\dots,m$$\n",
    "\n",
    "因为要产生target sequence的t时刻的输出，需要接收来自三个方向的信息，所以上面的公式很好理解。分别为：\n",
    "\n",
    "* $\\boldsymbol{s}_{t-1}$，前一个时刻的状态\n",
    "* $y_{t_1}$，前一个时刻的输入序列\n",
    "* $\\mathbf{c}_t$，当前时刻的context vector（上下文）\n",
    "\n",
    "其中，$\\mathbf{c}_t$就是`context vector`，表示`alignment scores`和source sequence的`hiddens states`的加权和：\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "\\mathbf{c}_t &= \\sum_{i=1}^n \\alpha_{t,i} \\boldsymbol{h}_i & \\small{\\text{; Context vector for output }y_t}\\\\\n",
    "\\alpha_{t,i} &= \\text{align}(y_t, x_i) & \\small{\\text{; How well two words }y_t\\text{ and }x_i\\text{ are aligned.}}\\\\\n",
    "&= \\frac{\\text{score}(\\boldsymbol{s}_{t-1}, \\boldsymbol{h}_i)}{\\sum_{i'=1}^n \\text{score}(\\boldsymbol{s}_{t-1}, \\boldsymbol{h}_{i'})} & \\small{\\text{; Softmax of some predefined alignment score.}}.\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "也就是说：\n",
    "\n",
    "$$\\mathbf{c}_t = \\sum_{i=1}^n \\frac{\\text{score}(\\boldsymbol{s}_{t-1}, \\boldsymbol{h}_i)}{\\sum_{i'=1}^n \\text{score}(\\boldsymbol{s}_{t-1}, \\boldsymbol{h}_{i'})} \\boldsymbol{h}_i$$\n",
    "\n",
    "根据公式可以看到：**$t$时刻的`context vector`需要计算所有的输入序列的隐状态**!\n",
    "\n",
    "$\\alpha_{t,i}$表示的是第$t$个时刻的输入和第$i$个输出之间的关系，那么所有的$t$和$i$组成的$\\alpha$就是**整个输入序列和输出序列之间的关系**！\n",
    "\n",
    "所以，这个$\\alpha$是一个二维矩阵！一般我们称呼它为`alignment matrix`！下面有一张示意图：\n",
    "\n",
    "<img src=\"images/attention_alignment_matrix.png\" width=300 height=300>\n",
    "\n",
    "\n",
    "那么，输入序列和输出序列的各个时间步之间的关系有多强烈呢？这就要用上面公式中的`score`函数来计算。\n",
    "\n",
    "`score`函数有很多选择，以下是常用的几种：\n",
    "\n",
    "* $\\text{score}(\\boldsymbol{s}_t, \\boldsymbol{h}_i) = \\mathbf{v}_a^\\top \\tanh(\\mathbf{W}_a[\\boldsymbol{s}_t; \\boldsymbol{h}_i])$，加性注意力（additive attention）\n",
    "* $\\text{score}(\\boldsymbol{s}_t, \\boldsymbol{h}_i) = \\boldsymbol{s}_t^\\top\\mathbf{W}_a\\boldsymbol{h}_i$，通用注意力（General）\n",
    "* $\\text{score}(\\boldsymbol{s}_t, \\boldsymbol{h}_i) = \\boldsymbol{s}_t^\\top\\boldsymbol{h}_i$，点积注意力（dot-product attention），属于乘性注意力（multiplicative）\n",
    "* $\\text{score}(\\boldsymbol{s}_t, \\boldsymbol{h}_i) = \\frac{\\boldsymbol{s}_t^\\top\\boldsymbol{h}_i}{\\sqrt{n}}$，上面的缩放版本\n",
    "\n",
    "上面的所有$\\text{W}_\\alpha$和$\\mathbf{v}_a^\\top$都是可训练的权值矩阵。\n",
    "\n",
    "[《Attention?Attention!》](https://lilianweng.github.io/lil-log/2018/06/24/attention-attention.html#a-family-of-attention-mechanisms)一文有一张比较全面的表格，展示了多种attention机制：\n",
    "![attention_mechanism_family](images/attention_mechanism_table.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tensorflow中的Attention机制\n",
    "\n",
    "`AttentionMechanism`是一个抽象类，或者说是接口。\n",
    "\n",
    "它的代码定义如下：\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class AttentionMechanism(object):\n",
    "\n",
    "  @property\n",
    "  def alignments_size(self):\n",
    "    raise NotImplementedError\n",
    "\n",
    "  @property\n",
    "  def state_size(self):\n",
    "    raise NotImplementedError"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "从这个接口定义，我们可以看出，attention机制只关注两点：\n",
    "\n",
    "* alignments_size\n",
    "* state_size\n",
    "\n",
    "在这里，我们看不出这两个属性是什么意思。那么我们先看看它的子类`_BaseAttentionMechanism`吧。\n",
    "\n",
    "根据文档，`_BazeAttentionMechanism`是一个最基本的注意力机制，它提供几个通用的功能，如下：\n",
    "\n",
    "* 存储query_layer和memory_layer\n",
    "* 预处理并存储memory\n",
    "\n",
    "什么是`query_layer`和`memory_layer`呢？我们留到后面解答。\n",
    "\n",
    "首先看看`_BaseAttentionMechanism`的`__init__`函数，它的签名如下：\n",
    "\n",
    "|参数名                   |解释                                           |\n",
    "|:-----------------------|:---------------------------------------------|\n",
    "|query_layer             |`tf.layers.Layer`的子类                        |\n",
    "|memory                  |通常是RNN encoder的输出                          |\n",
    "|probability_fn          |概率计算函数，默认是`softmax`                     |\n",
    "|memory_sequence_length  |序列长度                                        |\n",
    "|memory_layer            |`tf.layers.Layer`的子类                         |\n",
    "|check_inner_dims_defined|是否检查内部维度                                 |\n",
    "|score_mask_value        |掩码值，默认是负无穷                              |\n",
    "|name                    |变量域                                          |\n",
    "\n",
    "下面请看代码的注释，来理解初始化过程。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def __init__(self,\n",
    "               query_layer,\n",
    "               memory,\n",
    "               probability_fn,\n",
    "               memory_sequence_length=None,\n",
    "               memory_layer=None,\n",
    "               check_inner_dims_defined=True,\n",
    "               score_mask_value=None,\n",
    "               name=None):\n",
    "\n",
    "    if (query_layer is not None　and not isinstance(query_layer, layers_base.Layer)):\n",
    "      raise TypeError(\"query_layer is not a Layer: %s\" % type(query_layer).__name__)\n",
    "    if (memory_layer is not None　and not isinstance(memory_layer, layers_base.Layer)):\n",
    "      raise TypeError(\"memory_layer is not a Layer: %s\" % type(memory_layer).__name__)\n",
    "    \n",
    "    self._query_layer = query_layer\n",
    "    self._memory_layer = memory_layer\n",
    "    self.dtype = memory_layer.dtype\n",
    "    if not callable(probability_fn):\n",
    "      raise TypeError(\"probability_fn must be callable, saw type: %s\" %　type(probability_fn).__name__)\n",
    "    \n",
    "    # 默认是负无穷\n",
    "    if score_mask_value is None:\n",
    "      score_mask_value = dtypes.as_dtype(self._memory_layer.dtype).as_numpy_dtype(-np.inf)\n",
    "    \n",
    "    self._probability_fn = lambda score, prev: (  # pylint:disable=g-long-lambda\n",
    "        # _maybe_mask_score稍后解释\n",
    "        probability_fn(_maybe_mask_score(score, memory_sequence_length, score_mask_value),prev))\n",
    "    \n",
    "    with ops.name_scope(\n",
    "        name, \"BaseAttentionMechanismInit\", nest.flatten(memory)):\n",
    "      # 稍后解释\n",
    "      self._values = _prepare_memory(memory, memory_sequence_length,\n",
    "                                     check_inner_dims_defined=check_inner_dims_defined)\n",
    "      # shape = [B, T,...]\n",
    "      self._keys = self.memory_layer(self._values) if self.memory_layer else self._values\n",
    "      # 获取keys的第一个维度，即batch size\n",
    "      self._batch_size = self._keys.shape[0].value or array_ops.shape(self._keys)[0]\n",
    "      # 获取keys的第二个维度，即alignments size，这就是序列的最大长度\n",
    "      self._alignments_size = self._keys.shape[1].value or array_ops.shape(self._keys)[1]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `_maybe_mask_score`\n",
    "\n",
    "输入的张量score的形状是[B,T,..]，并且这个张量是对齐的，也就是每个序列的T这个维度的值都是一样的。但是实际上，我们同一批序列中，序列的长短是不一定一样的，也就是说我们的score张量中的某些值是无效的。我们需要把这些无效的值的位置找出来，然后给这些位置设置一个特定的数值，使得在这些位置的影响很小，趋于0。一般来说，设置一个**负无穷**。根据原来的序列的真实长度张量，以及我们的score张量，这样就可以找出哪些位置是无效的。当然，这要结合`probability_fn`函数，因为只有根据这个函数的性质，我们才知道赋予什么值，会使得这个函数的值趋于0。\n",
    "\n",
    "这也就说明了：如果我们的输入序列是已经对齐了的，那么我们可以直接设置`memory_sequence_length=None`，这样不就需要\n",
    "进一步处理了！\n",
    "\n",
    "这个函数的签名解释如下：\n",
    "\n",
    "|参数名                 |解释                                                |\n",
    "|:---------------------|:---------------------------------------------------|\n",
    "|score                 |分数张量，也就是我们上面提到的`score`函数的结果，形状是[B,T]|\n",
    "|memory_sequence_length|输入数据的长度张量，形状是[B]                           |\n",
    "|score_mask_value      |掩码值，默认是负无穷                                   |\n",
    "\n",
    "\n",
    "这个函数的实现如下：\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _maybe_mask_score(score, memory_sequence_length, score_mask_value):\n",
    "  # 如果序列长度未指定，则不进行掩码计算\n",
    "  if memory_sequence_length is None:\n",
    "    return score\n",
    "  message = (\"All values in memory_sequence_length must greater than zero.\")\n",
    "  with ops.control_dependencies(\n",
    "      [check_ops.assert_positive(memory_sequence_length, message=message)]):\n",
    "    # 进行sequence mask，得到的是一个Boolean类型（或者只有0和1）的张量，score元素为0的地方为False，其他位置为True\n",
    "    score_mask = array_ops.sequence_mask(memory_sequence_length, maxlen=array_ops.shape(score)[1])\n",
    "    # mask张量\n",
    "    score_mask_values = score_mask_value * array_ops.ones_like(score)\n",
    "    # 返回mask之后的张量，原来score位置为0的地方替换成了score_mask_value\n",
    "    return array_ops.where(score_mask, score, score_mask_values)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `_prepare_memory`\n",
    "\n",
    "这个函数是把输入张量转化成掩码之后的张量。为什么这里需要进行掩码处理呢？\n",
    "\n",
    "首先memory张量的形状是[B,T,...]，其中T是这一批次的数据中的最大长度。为了搞清楚这T个time step之间，哪些是有效位置，哪些是无效的（通常进行对齐添加的标记），我们需要把这些无效的位置找出来，然后数值上设为0。这样这些无效位置就不会影响后面的计算。\n",
    "\n",
    "当然，如果你手动对齐了输入序列（在较短的序列后面补0），保证每个序列的长度都一样，那么你可以设置参数`memory_sequence_length=None`，这样就不需要额外处理了。同样的，即使你对齐处理过，你也可以传入一个`memory_sequence_length`张量，这样不会有影响。\n",
    "\n",
    "签名如下：\n",
    "\n",
    "|参数名                   |解释                           |\n",
    "|:-----------------------|:-----------------------------|\n",
    "|memory                  |输入张量，形状为[B,T,...]        |\n",
    "|memory_sequence_length  |序列长度张量，形状为[B]           |\n",
    "|check_inner_dims_defined|Boolean，检查最外面两个维度是否定义|\n",
    "\n",
    "注意，掩码只是作用于[B,T]两个维度，之后如果还有其他维度，应该处理使之不进行掩码计算。这个在代码中有很明显的处理。\n",
    "\n",
    "下面，详细解释代码：\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def _prepare_memory(memory, memory_sequence_length, check_inner_dims_defined):\n",
    "    memory = nest.map_structure(\n",
    "        lambda m: ops.convert_to_tensor(m, name=\"memory\"), memory)\n",
    "    if memory_sequence_length is not None:\n",
    "        memory_sequence_length = ops.convert_to_tensor(memory_sequence_length,\n",
    "                                                       name=\"memory_sequence_length\")\n",
    "\n",
    "    # 检查innder dim\n",
    "    if check_inner_dims_defined:\n",
    "        def _check_dims(m):\n",
    "            # 检查第三个维度及其之后的维度是否存在\n",
    "            if not m.get_shape()[2:].is_fully_defined():\n",
    "                raise ValueError(\n",
    "                    \"Expected memory %s to have fully defined inner dims, \"\n",
    "                    \"but saw shape: %s\" % (m.name, m.get_shape()))\n",
    "            nest.map_structure(_check_dims, memory)\n",
    "\n",
    "    if memory_sequence_length is None:\n",
    "        seq_len_mask = None\n",
    "    else:\n",
    "        # 进行sequence_mask\n",
    "        seq_len_mask = array_ops.sequence_mask(\n",
    "            memory_sequence_length,\n",
    "            maxlen=array_ops.shape(nest.flatten(memory)[0])[1],\n",
    "            dtype=nest.flatten(memory)[0].dtype)\n",
    "        # 获取batch size\n",
    "        seq_len_batch_size = (memory_sequence_length.shape[0].value or\n",
    "                              array_ops.shape(memory_sequence_length)[0])\n",
    "\n",
    "    def _maybe_mask(m, seq_len_mask):\n",
    "        # 获取维度数量，也就是阶\n",
    "        rank = m.get_shape().ndims\n",
    "        rank = rank if rank is not None else array_ops.rank(m)\n",
    "        # 取出B,T之外的维度，设置成1，当计算掩码的时候，这些维度的值会保留原来的值\n",
    "        extra_ones = array_ops.ones(rank - 2, dtype=dtypes.int32)\n",
    "        # 获取batch size\n",
    "        m_batch_size = m.shape[0].value or array_ops.shape(m)[0]\n",
    "        if memory_sequence_length is not None:\n",
    "            message = (\n",
    "                \"memory_sequence_length and memory tensor batch sizes do not match.\")\n",
    "            with ops.control_dependencies([\n",
    "                check_ops.assert_equal(seq_len_batch_size, m_batch_size,\n",
    "                                       message=message)]):\n",
    "                # 将seq_len_mask重塑成[B,T,...]形状的张量\n",
    "                seq_len_mask = array_ops.reshape(\n",
    "                    seq_len_mask,\n",
    "                    array_ops.concat((array_ops.shape(seq_len_mask), extra_ones), 0))\n",
    "                # 掩码计算，并且返回结果\n",
    "                return m * seq_len_mask\n",
    "        else:\n",
    "            return m\n",
    "\n",
    "    return nest.map_structure(lambda m: _maybe_mask(m, seq_len_mask), memory)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "以上首先就可以回答`AttentionMechanism`类的第一个函数是什么这个问题。因为`_BaseAttentionMechansim`的实现很简单：\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "  @property\n",
    "  def alignments_size(self):\n",
    "    # __init__函数里面，self._alignments_size = self._keys.shape[1].value or array_ops.shape(self._keys)[1]\n",
    "    # 实际上就是这一批序列的最大长度\n",
    "    return self._alignments_size"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "那么第二个问题，`state_size`又是什么呢？在`_BaseAttentionMechanism`里面的实现，和上面一样：\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "  @property\n",
    "  def state_size(self):\n",
    "    return self._alignments_size"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "也就是说，默认情况下，`alignments_size`和`state_size`是一样的！\n",
    "\n",
    "那么为什么这两个值默认是一样的呢？\n",
    "\n",
    "回顾文章开始的attention机制的定义，我们发现知道，`alignment matrix`的两个维度的大小，分别是输入序列和输出序列的长度，并且，在编解码的过程中，每一个时间步都会产生一个`hidden state`，所以`hidden state`的数量也是和序列长度一样的。\n",
    "\n",
    "而我们这里的`alinments_size`就是对应`alignment matrix`，`state`就是对应`hidden state`的！\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Query，Keys和Values\n",
    "\n",
    "你可能已经发现了，在`_BaseAttentionMechanism`的`__init__`函数里，出现了`self._query_layer`、`self._memory_layer`、`self._keys`和`self._values`等成员变量。那么这些变量指的是什么呢？\n",
    "\n",
    "[Attention is all you need](https://arxiv.org/pdf/1706.03762.pdf)一文有一个对attention机制的解释：\n",
    "\n",
    "> An attention function can be described as mapping a query and a set of key-values pairs to an output, where the query, keys, values and output are all vectors. The output is computed as a weighted sum of the values, where the weight assigned to each value is computed by a compatibility function of the query with the corresponding key.\n",
    "\n",
    "翻译一下就是：\n",
    "\n",
    "> 注意力函数可以描述为一个映射过程：把query和key-value键值对映射到一个输出。其中query，keys, values和output都是矢量。output是values的加权和，其中每一个value的权重都是通过一个关于key和query的兼容性函数计算得来的。\n",
    "\n",
    "我再加一点解释：\n",
    "\n",
    "* 所谓的关于key和query的兼容性函数，就是我们的`score`函数！\n",
    "\n",
    "所以关于attention机制，可以再精简一点，概括为：**Attention机制就是，从key-value键值对中，通过确定query和key的相似度，然后再通过value获得最终输出的映射过程。**\n",
    "\n",
    "你已经发现了，上面的定义都是围绕`keys`、`query`和`values`进行的！！！\n",
    "\n",
    "那么言归正传，什么是`keys`、`values`和`query`呢？\n",
    "\n",
    "在`_BaseAttentionMechanism`里面，`keys`和`values`的关系是这样的：\n",
    "\n",
    "```python\n",
    "self._values = _prepare_memory(\n",
    "          memory, memory_sequence_length,\n",
    "          check_inner_dims_defined=check_inner_dims_defined)\n",
    "self._keys = (\n",
    "          self.memory_layer(self._values) if self.memory_layer  # pylint: disable=not-callable\n",
    "          else self._values)\n",
    "```\n",
    "`_prepare_memory`上面已经解释过了。所以：\n",
    "\n",
    "* `values`就是`memory或者masked memory`！（memory一般是RNN encoder的输出`encoder_output`！）\n",
    "* `keys`就是`values`或者`memory_layer处理过的values`！(keys和values可以认为是一样的！)\n",
    "\n",
    "那么，这里为什么没有出现`query`呢？很简单，因为`query`是在调用的时候传入的参数，所以这里没有`query`！具体的可以从`_BaseAttentionMechanism`的子类的`__call__`函数解读出来。\n",
    "\n",
    "总结一下：\n",
    "\n",
    "* `values`就是`encoder_output`或者`masked encoder_output`\n",
    "* `keys`就是`values`或者是经过`memory_layer`处理后的`values`\n",
    "* `query`是调用的时候传入的参数，一般与来自RNN decoder\n",
    "\n",
    "疑问: **这里的`query`、`keys`和`values`怎么和`score`函数对应起来？**\n",
    "（提示，可以参考attention_wrapper代码。）\n",
    "\n",
    "这个问题留到后面解释。\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### memory_layer和query_layer\n",
    "\n",
    "上面还有两个疑惑，`memory_layer`和`query_layer`是什么？\n",
    "\n",
    "`_BaseAttentionMechanism`没有解答，但是它的子类`LuongAttention`和`BahdanauAttention`可以找到答案。\n",
    "\n",
    "`LuongAttention`的`__init__`函数有定义：\n",
    "\n",
    "```python\n",
    "query_layer=None,\n",
    "memory_layer=layers_core.Dense(\n",
    "            num_units, name=\"memory_layer\", use_bias=False, dtype=dtype)\n",
    "```\n",
    "\n",
    "在`LuongAttention`里，`query_layer`是`None`，`memory_layer`只是一个**全连接层**！，并且激活函数默认是`None`，也就是说**memory_layer只是对memory做一个简单的线性变换！**\n",
    "\n",
    "`BahdanauAttention`是这样定义的：\n",
    "\n",
    "```python\n",
    "query_layer=layers_core.Dense(\n",
    "            num_units, name=\"query_layer\", use_bias=False, dtype=dtype)\n",
    "memory_layer=layers_core.Dense(\n",
    "            num_units, name=\"memory_layer\", use_bias=False, dtype=dtype)\n",
    "```\n",
    "\n",
    "在`BahdanauAttention`里，`query_layer`和`memory_layer`都是一个全连接层，激活函数默认也是`None`，所以`query_layer`和`memory_layer`分别是对`query`和`memory`的线性变换。\n",
    "\n",
    "现在`memory_layer`和`query_layer`已经清楚了，但是现在又有一个新的问题：**`LuongAttention`和`BahdanauAttention`有什么区别呢？**\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### LuongAttention和BahdanauAttention\n",
    "\n",
    "不同的attention机制，最大和最关键的区别就在于它们各自的`score`函数。\n",
    "\n",
    "我们首先来看看`LuongAttention`的`score`函数：\n",
    "\n",
    "```python\n",
    "def _luong_score(query, keys, scale):\n",
    "  # query shape: [batch_size, depth]\n",
    "  # keys shape: [batch_size, max_time, depth]\n",
    "  depth = query.get_shape()[-1]\n",
    "  key_units = keys.get_shape()[-1]\n",
    "  if depth != key_units:\n",
    "    raise ValueError(\n",
    "        \"Incompatible or unknown inner dimensions between query and keys.  \"\n",
    "        \"Query (%s) has units: %s.  Keys (%s) have units: %s.  \"\n",
    "        \"Perhaps you need to set num_units to the keys' dimension (%s)?\"\n",
    "        % (query, depth, keys, key_units, key_units))\n",
    "  dtype = query.dtype\n",
    "\n",
    "  # 为后面的矩阵乘法做准备，在第二个维度上扩充一个维度\n",
    "  # [batch_size, depth] -> [batch_size, 1, depth]\n",
    "  query = array_ops.expand_dims(query, 1)\n",
    "\n",
    "  # [batch_size, 1, depth]*[batch_size, depth, max_time] -> [batch_size, 1, max_time]\n",
    "  score = math_ops.matmul(query, keys, transpose_b=True)\n",
    "  # 将第二个维度压缩掉，[batch_size, 1, max_time] -> [batch_size, max_time]\n",
    "  score = array_ops.squeeze(score, [1])\n",
    "  \n",
    "  # 缩放，可选项\n",
    "  if scale:\n",
    "    # Scalar used in weight scaling\n",
    "    g = variable_scope.get_variable(\n",
    "        \"attention_g\", dtype=dtype,\n",
    "        initializer=init_ops.ones_initializer, shape=())\n",
    "    score = g * score\n",
    "  return score\n",
    "```\n",
    "\n",
    "`LuongAttention`计算`socre`的过程很简单：**query和keys做矩阵乘法而已！**\n",
    "\n",
    "从数学上来说，这里的`score`计算公式如下：\n",
    "\n",
    "$$\\text{score}(\\boldsymbol{s}_t, \\boldsymbol{h}_i) = \\boldsymbol{s}_t^\\top\\boldsymbol{h}_i$$\n",
    "\n",
    "所以，`LuongAttention`属于**multiplicative attention**！，具体一点，就是**dot-product attention**!\n",
    "\n",
    "我们继续看看`BahdanauAttention`的`score`函数：\n",
    "\n",
    "```python\n",
    "def _bahdanau_score(processed_query, keys, normalize):\n",
    "  # processed_query shape: [batch_size, depth]\n",
    "  # keys shape: [batch_size, max_time, depth]\n",
    "  dtype = processed_query.dtype\n",
    "  # 获取depth(num_units)，是第三个维度\n",
    "  num_units = keys.shape[2].value or array_ops.shape(keys)[2]\n",
    "  # 将query变形，[batch_size, depth] -> [batch_size, 1, depth]\n",
    "  processed_query = array_ops.expand_dims(processed_query, 1)\n",
    "  v = variable_scope.get_variable(\n",
    "      \"attention_v\", [num_units], dtype=dtype)\n",
    "  # 归一化，可选项\n",
    "  if normalize:\n",
    "    # Scalar used in weight normalization\n",
    "    g = variable_scope.get_variable(\n",
    "        \"attention_g\", dtype=dtype,\n",
    "        initializer=init_ops.constant_initializer(math.sqrt((1. / num_units))),\n",
    "        shape=())\n",
    "    # tanh激活函数的偏置\n",
    "    b = variable_scope.get_variable(\n",
    "        \"attention_b\", [num_units], dtype=dtype,\n",
    "        initializer=init_ops.zeros_initializer())\n",
    "    # normed_v = g * v / ||v||\n",
    "    normed_v = g * v * math_ops.rsqrt(\n",
    "        math_ops.reduce_sum(math_ops.square(v)))\n",
    "    # query和keys直接相加\n",
    "    return math_ops.reduce_sum(\n",
    "        normed_v * math_ops.tanh(keys + processed_query + b), [2])\n",
    "  else:\n",
    "    # query和keys直接相加\n",
    "    return math_ops.reduce_sum(v * math_ops.tanh(keys + processed_query), [2])\n",
    "```\n",
    "\n",
    "可以看出，`BahdanauAttention`属于**additive attention**！\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 参考文章\n",
    "\n",
    "1.[Attention?Attention!](https://lilianweng.github.io/lil-log/2018/06/24/attention-attention.html#a-family-of-attention-mechanisms)  \n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
